package com.pai4j.aigc.llm.service;

import com.alibaba.fastjson2.JSONObject;
import com.pai4j.aigc.llm.entity.LLMModelEntity;
import org.springframework.stereotype.Service;

@Service
public class PricingService {

    /**
     * 计算请求的费用（以美分为单位）
     *
     * @param model 包含定价参数的模型实体，其parameters字段应包含JSON格式的定价信息
     * @param promptTokens 输入token数量
     * @param completionTokens 输出token数量
     * @return 计算出的总费用（美分），如果无法计算则返回null
     */
    public Integer calcCostCents(LLMModelEntity model, Integer promptTokens, Integer completionTokens) {
        // 检查模型和参数是否有效
        if (model == null || model.getParameters() == null) return null;
        try {
            // 解析模型参数JSON
            JSONObject obj = JSONObject.parseObject(model.getParameters());
            if (obj == null) return null;

            // 提取输入和输出的单价参数
            Integer inPer1k = obj.getInteger("input_cost_per_1k_cents");
            Integer outPer1k = obj.getInteger("output_cost_per_1k_cents");

            // 如果两个参数都不存在，则无法计算费用
            if (inPer1k == null && outPer1k == null) return null;

            // 分别计算输入和输出费用
            double inCost = 0;
            double outCost = 0;
            if (inPer1k != null && promptTokens != null) {
                inCost = (promptTokens / 1000.0) * inPer1k;
            }
            if (outPer1k != null && completionTokens != null) {
                outCost = (completionTokens / 1000.0) * outPer1k;
            }

            // 计算总费用并进行边界处理
            long cents = Math.round(inCost + outCost);
            if (cents < 0) cents = 0;
            if (cents > Integer.MAX_VALUE) cents = Integer.MAX_VALUE;
            return (int) cents;
        } catch (Exception e) {
            return null;
        }
    }
}

